Proposal for Caching the Within-Preconditioner#1264
Proposal for Caching the Within-Preconditioner#1264
Conversation
…r_ attribute for reuse across model fits
…en inner tolerance is tighened. in addition, add inner tolerance tighening to fepois
leostimpfle
left a comment
There was a problem hiding this comment.
This looks very nice! After an initial look (mostly at the design), I have three (fairly minor) comments:
- I would favour more consistent use of
typing.Literalto avoid passingstrarguments (e.g.,preconditioner_type) - Many of the
object.__setattr__inside the frozen dataclasses can probably be avoided so that__post_init__just calls validations (but does not modify provided arguments) - I'm wondering if adding a
demeanmethod toBaseDemeanerwould be nicer than the explicitdispatch_demean
| n_obs: int, | ||
| n_factors: int, | ||
| factor_cardinalities: tuple[int, ...], | ||
| preconditioner_type: str, |
There was a problem hiding this comment.
Can we use a typing.Literal for clarity here and anywhere preconditioner_type is used? For example, PreconditionerType = Literal["additive", "multiplicative"]
| class WithinDemeaner(BaseDemeaner): | ||
| """Krylov-based demeaner configuration for the Rust `within` backend.""" | ||
|
|
||
| krylov_method: str = "cg" |
There was a problem hiding this comment.
Also use a typing.Literal, e.g., Literal["cg", "gmres"]?
| class LsmrDemeaner(BaseDemeaner): | ||
| """Sparse LSMR demeaner for CPU and GPU backends.""" | ||
|
|
||
| precision: str = "float64" |
There was a problem hiding this comment.
Use Literal["float32", "float64"]?
| krylov_method: str, | ||
| preconditioner_type: str, | ||
| ) -> tuple[str, str]: | ||
| krylov_method = krylov_method.lower() |
There was a problem hiding this comment.
Do we need this kind of string parsing or should we tighten the arguments to a Literal (see other comments)?
| maxiter: int = 1_000, | ||
| krylov_method: str = "cg", | ||
| gmres_restart: int = 30, | ||
| preconditioner_type: str = "additive", |
There was a problem hiding this comment.
Tighten to Literal["additive", "multiplicative"]?
| object.__setattr__( | ||
| self, | ||
| "fixef_tol", | ||
| _validate_positive_float(self.fixef_tol, "fixef_tol"), | ||
| ) | ||
| object.__setattr__( | ||
| self, | ||
| "fixef_maxiter", | ||
| _validate_positive_int(self.fixef_maxiter, "fixef_maxiter"), | ||
| ) |
There was a problem hiding this comment.
Why not just call _validate_positive_float and _validate_positive_int instead of reassigning the attributes?
| backend = self.backend.lower() | ||
| if backend not in {"numba", "rust", "jax"}: | ||
| raise ValueError("`backend` must be one of 'numba', 'rust', or 'jax'.") | ||
| object.__setattr__(self, "backend", backend) |
There was a problem hiding this comment.
We can avoid the reassign when tightening backend to be case sensitive (consistent with the type hint)?
| ) | ||
| from pyfixest.estimation.internals.literals import DemeanerBackendOptions | ||
|
|
||
| ResolvedDemeaner: TypeAlias = AnyDemeaner |
There was a problem hiding this comment.
Why is this introduced? Couldn't we just type hint with AnyDemeaner?
| return replace(demeaner, fixef_tol=tol) | ||
|
|
||
|
|
||
| def dispatch_demean( |
There was a problem hiding this comment.
How about giving demeaner a method demean that avoids this dispatch function?
Summary
This PR introduces a typed fixed-effects demeaning API, passes it through the code base, and adds reusable
withinpreconditioners to speed up repeated multi-way FE estimation.The main user-facing additions are:
demeaner=objects:MapDemeaner(...)WithinDemeaner(...)LsmrDemeaner(...)demeanernow takes precedence overdemeaner_backend,fixef_tol, andfixef_maxiterfeolsnow exposesfit.preconditioner_forWithinDemeanerWithinDemeaner(preconditioner=...)WithinPreconditionercan be pickled across sessions, following the upstreamwithin-pypatternWhat Changed
1. Typed demeaner API
We now support a typed
demeaner=interface instead of relying only on string backend selection viademeaner_backend. This makes the demeaner backend configuration explicit, extensible, and easier to validate.The typed demeaner objects control actual runtime execution.
WithinDemeaneroptions flowing into thewithinbackendLsmrDemeaneroptions flowing into the SciPy/CuPy solversMapDemeaneroptions flowing into the MAP solvers2. Reusable
withinpreconditionersFor multi-way FE
WithinDemeanerfits, we now build and cache a reusableWithinPreconditioner.That preconditioner is:
feolsasfit.preconditioner_WithinDemeaner(preconditioner=...)feglmandfepois3. IWLS reuse
For
feglmandfepois, we now reuse a fit-localwithinpreconditioner across IWLS iterations.Current policy is intentionally simple:
@schroedk - For generalized linear models, we need to demean repeatedly with "different weights". Here we currently say it is better to run on "stale" preconditions for a while than to update the preconditioner in every iteration. This can of course be refined later.
4. Persistence
WithinPreconditioneris pickleable across sessions.This mirrors
within-pyclosely:FePreconditioneris serialized withpostcard__reduce__Key Design Choices
Why expose preconditioners, but not solvers?
We intentionally expose preconditioners as reusable objects, but not solvers.
Because reusing a
withinpreconditioner is where most of the practical speedup comes from. It is "ok" if the preconditioner becomes "a bit stale" - this is a feature, not a bug imo. Plus, we do not need to check that the fixed effects and the sample and weights are identical, which is what we would have needed to do if we allowed users to pass solvers.For iterated weighted least squares as used for GLMs, the preconditioner is the better choice as in each iteration, a stale preconditioner will be "good enough".
Stay close to
within-pyTried to stay close to the implementation in
within-py.In particular:
within-pyWorkflow
Typical usage now looks like:
For persistent reuse, you can do:
And then fit the regression model.